// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <melon/cportability.h>

#include <pollux/expression/expr.h>
#include <pollux/expression/vector_function.h>
#include <pollux/type/floating_point_util.h>
#include <pollux/vector/decoded_vector.h>

namespace kumo::pollux::functions {
    namespace {
        template<typename T>
        inline bool isPrimitiveEqual(const T &lhs, const T &rhs) {
            if constexpr (std::is_floating_point_v<T>) {
                return util::floating_point::NaNAwareEquals<T>{}(lhs, rhs);
            } else {
                return lhs == rhs;
            }
        }

        // Find the index of the first match for primitive types.
        template<
            bool useCustomComparison,
            TypeKind kind,
            typename std::enable_if_t<
                !useCustomComparison && TypeTraits<kind>::isPrimitiveType,
                int>  = 0>
        void applyTypedFirstMatch(
            const SelectivityVector &rows,
            DecodedVector &arrayDecoded,
            const DecodedVector &elementsDecoded,
            const DecodedVector &searchDecoded,
            FlatVector<int64_t> &flatResult) {
            using T = typename TypeTraits<kind>::NativeType;

            auto baseArray = arrayDecoded.base()->as<ArrayVector>();
            auto rawSizes = baseArray->rawSizes();
            auto rawOffsets = baseArray->rawOffsets();
            auto indices = arrayDecoded.indices();

            if (elementsDecoded.isIdentityMapping() && !elementsDecoded.may_have_nulls() &&
                searchDecoded.isConstantMapping()) {
                // Fast path for array vector of boolean.
                if constexpr (std::is_same_v<bool, T>) {
                    auto rawElements = elementsDecoded.data<uint64_t>();

                    auto search = bits::isBitSet(searchDecoded.data<uint64_t>(), 0);

                    rows.applyToSelected([&](auto row) {
                        auto size = rawSizes[indices[row]];
                        auto offset = rawOffsets[indices[row]];

                        int i;
                        for (i = 0; i < size; i++) {
                            if (bits::isBitSet(rawElements, offset + i) == search) {
                                flatResult.set(row, i + 1);
                                break;
                            }
                        }
                        if (i == size) {
                            flatResult.set(row, 0);
                        }
                    });
                    return;
                }

                // Fast path for array vector of types other than boolean.
                auto rawElements = elementsDecoded.data<T>();

                auto search = searchDecoded.value_at<T>(0);

                rows.applyToSelected([&](auto row) {
                    auto size = rawSizes[indices[row]];
                    auto offset = rawOffsets[indices[row]];

                    int i;
                    for (i = 0; i < size; i++) {
                        if (isPrimitiveEqual<T>(rawElements[offset + i], search)) {
                            flatResult.set(row, i + 1);
                            break;
                        }
                    }
                    if (i == size) {
                        flatResult.set(row, 0);
                    }
                });
                return;
            }

            // Regular path where no assumption is made about the encodings of
            // searchDecoded and elementsDecoded.
            rows.applyToSelected([&](auto row) {
                auto size = rawSizes[indices[row]];
                auto offset = rawOffsets[indices[row]];

                auto search = searchDecoded.value_at<T>(row);

                int i;
                for (i = 0; i < size; i++) {
                    if (!elementsDecoded.is_null_at(offset + i) &&
                        isPrimitiveEqual<T>(elementsDecoded.value_at<T>(offset + i), search)) {
                        flatResult.set(row, i + 1);
                        break;
                    }
                }
                if (i == size) {
                    flatResult.set(row, 0);
                }
            });
        }

        // Find the index of the first match for complex types.
        template<
            bool useCustomComparison,
            TypeKind kind,
            typename std::enable_if_t<
                useCustomComparison || !TypeTraits<kind>::isPrimitiveType,
                int>  = 0>
        void applyTypedFirstMatch(
            const SelectivityVector &rows,
            DecodedVector &arrayDecoded,
            DecodedVector &elementsDecoded,
            DecodedVector &searchDecoded,
            FlatVector<int64_t> &flatResult) {
            auto baseArray = arrayDecoded.base()->as<ArrayVector>();
            auto rawSizes = baseArray->rawSizes();
            auto rawOffsets = baseArray->rawOffsets();
            auto indices = arrayDecoded.indices();

            auto elementsBase = elementsDecoded.base();
            auto elementIndices = elementsDecoded.indices();

            auto searchBase = searchDecoded.base();
            auto searchIndices = searchDecoded.indices();

            rows.applyToSelected([&](auto row) {
                auto size = rawSizes[indices[row]];
                auto offset = rawOffsets[indices[row]];
                auto searchIndex = searchIndices[row];

                int i;
                for (i = 0; i < size; i++) {
                    if (!elementsDecoded.is_null_at(offset + i) &&
                        elementsBase->equal_value_at(
                            searchBase, elementIndices[offset + i], searchIndex)) {
                        flatResult.set(row, i + 1);
                        break;
                    }
                }
                if (i == size) {
                    flatResult.set(row, 0);
                }
            });
        }

        MELON_ALWAYS_INLINE void getLoopBoundary(
            const int *rawSizes,
            const int *indices,
            vector_size_t row,
            int64_t &instance,
            int &startIndex,
            int &endIndex,
            int &step) {
            auto size = rawSizes[indices[row]];

            startIndex = instance > 0 ? 0 : size - 1;
            endIndex = instance > 0 ? size : -1;
            step = instance > 0 ? 1 : -1;
            instance = std::abs(instance);
        }

        // Find the index of the instance-th match for primitive types.
        template<
            bool useCustomComparison,
            TypeKind kind,
            typename std::enable_if_t<
                !useCustomComparison && TypeTraits<kind>::isPrimitiveType,
                int>  = 0>
        void applyTypedWithInstance(
            const SelectivityVector &rows,
            exec::EvalCtx &context,
            DecodedVector &arrayDecoded,
            const DecodedVector &elementsDecoded,
            const DecodedVector &searchDecoded,
            const DecodedVector &instanceDecoded,
            FlatVector<int64_t> &flatResult) {
            using T = typename TypeTraits<kind>::NativeType;

            auto baseArray = arrayDecoded.base()->as<ArrayVector>();
            auto rawSizes = baseArray->rawSizes();
            auto rawOffsets = baseArray->rawOffsets();
            auto indices = arrayDecoded.indices();

            int startIndex;
            int endIndex;
            int step;

            if (elementsDecoded.isIdentityMapping() && !elementsDecoded.may_have_nulls() &&
                searchDecoded.isConstantMapping() &&
                instanceDecoded.isConstantMapping()) {
                const auto instance = instanceDecoded.value_at<int64_t>(0);

                try {
                    POLLUX_USER_CHECK_NE(
                        instance,
                        0,
                        "array_position cannot take a 0-valued instance argument.");
                } catch (...) {
                    context.setErrors(rows, std::current_exception());
                    return;
                }

                // Fast path for array vector of boolean.
                if constexpr (std::is_same_v<bool, T>) {
                    auto rawElements = elementsDecoded.data<uint64_t>();

                    auto search = bits::isBitSet(searchDecoded.data<uint64_t>(), 0);

                    rows.applyToSelected([&](auto row) {
                        auto offset = rawOffsets[indices[row]];
                        auto remaining = instance;
                        getLoopBoundary(
                            rawSizes, indices, row, remaining, startIndex, endIndex, step);

                        int i;
                        for (i = startIndex; i != endIndex; i += step) {
                            if (bits::isBitSet(rawElements, offset + i) == search) {
                                if (--remaining == 0) {
                                    flatResult.set(row, i + 1);
                                    break;
                                }
                            }
                        }
                        if (i == endIndex) {
                            flatResult.set(row, 0);
                        }
                    });
                    return;
                }

                // Fast path for array vector of types other than boolean.
                auto rawElements = elementsDecoded.data<T>();

                auto search = searchDecoded.value_at<T>(0);

                rows.applyToSelected([&](auto row) {
                    auto offset = rawOffsets[indices[row]];
                    auto remaining = instance;
                    getLoopBoundary(
                        rawSizes, indices, row, remaining, startIndex, endIndex, step);

                    int i;
                    for (i = startIndex; i != endIndex; i += step) {
                        if (isPrimitiveEqual<T>(rawElements[offset + i], search)) {
                            if (--remaining == 0) {
                                flatResult.set(row, i + 1);
                                break;
                            }
                        }
                    }
                    if (i == endIndex) {
                        flatResult.set(row, 0);
                    }
                });
                return;
            }

            // Regular path where no assumption is made about the encodings of
            // searchDecoded and elementsDecoded.
            context.applyToSelectedNoThrow(rows, [&](auto row) {
                auto offset = rawOffsets[indices[row]];
                auto search = searchDecoded.value_at<T>(row);

                auto instance = instanceDecoded.value_at<int64_t>(row);
                POLLUX_USER_CHECK_NE(
                    instance,
                    0,
                    "array_position cannot take a 0-valued instance argument.");

                getLoopBoundary(
                    rawSizes, indices, row, instance, startIndex, endIndex, step);

                int i;
                for (i = startIndex; i != endIndex; i += step) {
                    if (!elementsDecoded.is_null_at(offset + i) &&
                        isPrimitiveEqual<T>(elementsDecoded.value_at<T>(offset + i), search)) {
                        --instance;
                        if (instance == 0) {
                            flatResult.set(row, i + 1);
                            break;
                        }
                    }
                }
                if (i == endIndex) {
                    flatResult.set(row, 0);
                }
            });
        }

        // Find the index of the instance-th match for complex types.
        template<
            bool useCustomComparison,
            TypeKind kind,
            typename std::enable_if_t<
                useCustomComparison || !TypeTraits<kind>::isPrimitiveType,
                int>  = 0>
        void applyTypedWithInstance(
            const SelectivityVector &rows,
            exec::EvalCtx &context,
            DecodedVector &arrayDecoded,
            DecodedVector &elementsDecoded,
            DecodedVector &searchDecoded,
            const DecodedVector &instanceDecoded,
            FlatVector<int64_t> &flatResult) {
            auto baseArray = arrayDecoded.base()->as<ArrayVector>();
            auto rawSizes = baseArray->rawSizes();
            auto rawOffsets = baseArray->rawOffsets();
            auto indices = arrayDecoded.indices();

            auto elementsBase = elementsDecoded.base();
            auto elementIndices = elementsDecoded.indices();
            auto searchBase = searchDecoded.base();
            auto searchIndices = searchDecoded.indices();

            int startIndex;
            int endIndex;
            int step;

            context.applyToSelectedNoThrow(rows, [&](auto row) {
                auto offset = rawOffsets[indices[row]];
                auto searchIndex = searchIndices[row];

                auto instance = instanceDecoded.value_at<int64_t>(row);
                POLLUX_USER_CHECK_NE(
                    instance,
                    0,
                    "array_position cannot take a 0-valued instance argument.");

                getLoopBoundary(
                    rawSizes, indices, row, instance, startIndex, endIndex, step);

                int i;
                for (i = startIndex; i != endIndex; i += step) {
                    if (!elementsDecoded.is_null_at(offset + i) &&
                        elementsBase->equal_value_at(
                            searchBase, elementIndices[offset + i], searchIndex)) {
                        --instance;
                        if (instance == 0) {
                            flatResult.set(row, i + 1);
                            break;
                        }
                    }
                }
                if (i == endIndex) {
                    flatResult.set(row, 0);
                }
            });
        }

        template<bool useCustomComparison>
        void applyInternal(
            const SelectivityVector &rows,
            std::vector<VectorPtr> &args,
            exec::EvalCtx &context,
            VectorPtr &result) {
            context.ensure_writable(rows, BIGINT(), result);
            auto flatResult = result->as_flat_vector<int64_t>();

            exec::DecodedArgs decodedArgs(rows, args, context);
            auto elements = decodedArgs.at(0)->base()->as<ArrayVector>()->elements();
            exec::LocalSelectivityVector nestedRows(context, elements->size());
            nestedRows.get()->setAll();
            exec::LocalDecodedVector elementsHolder(
                context, *elements, *nestedRows.get());

            if (args.size() == 2) {
                POLLUX_DYNAMIC_TEMPLATE_TYPE_DISPATCH(
                    applyTypedFirstMatch,
                    useCustomComparison,
                    args[1]->type_kind(),
                    rows,
                    *decodedArgs.at(0),
                    *elementsHolder.get(),
                    *decodedArgs.at(1),
                    *flatResult);
            } else {
                const auto &instanceVector = args[2];
                POLLUX_CHECK(instanceVector->type()->is_bigint());

                POLLUX_DYNAMIC_TEMPLATE_TYPE_DISPATCH(
                    applyTypedWithInstance,
                    useCustomComparison,
                    args[1]->type_kind(),
                    rows,
                    context,
                    *decodedArgs.at(0),
                    *elementsHolder.get(),
                    *decodedArgs.at(1),
                    *decodedArgs.at(2),
                    *flatResult);
            }
        }

        class ArrayPositionFunction : public exec::VectorFunction {
        public:
            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr & /* outputType */,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                const auto &array_vector = args[0];
                const auto &searchVector = args[1];
                POLLUX_CHECK(array_vector->type()->is_array());
                POLLUX_CHECK(array_vector->type()->as_array().elementType()->kindEquals(
                    searchVector->type()));

                if (searchVector->type()->providesCustomComparison()) {
                    applyInternal<true>(rows, args, context, result);
                } else {
                    applyInternal<false>(rows, args, context, result);
                }
            }

            static std::vector<std::shared_ptr<exec::FunctionSignature> > signatures() {
                return {
                    // array(T), T -> int64_t
                    exec::FunctionSignatureBuilder()
                    .typeVariable("T")
                    .returnType("bigint")
                    .argumentType("array(T)")
                    .argumentType("T")
                    .build(),

                    // array(T), T, int64_t -> int64_t
                    exec::FunctionSignatureBuilder()
                    .typeVariable("T")
                    .returnType("bigint")
                    .argumentType("array(T)")
                    .argumentType("T")
                    .argumentType("bigint")
                    .build()
                };
            }
        };
    } // namespace

    POLLUX_DECLARE_VECTOR_FUNCTION(
        udf_array_position,
        ArrayPositionFunction::signatures(),
        std::make_unique<ArrayPositionFunction>());
} // namespace kumo::pollux::functions
